import yaml
import os

class ConfigLoader:
    def __init__(self, config_path=None):
        if config_path is None:
            # 使用相对于当前文件的路径
            current_dir = os.path.dirname(os.path.abspath(__file__))
            self.config_path = os.path.join(current_dir, "config.yaml")
        else:
            self.config_path = config_path
        self.config = self._load_config()

    def _load_config(self):
        """Load configuration from YAML file"""
        if not os.path.exists(self.config_path):
            return self._get_default_config()
        
        with open(self.config_path, 'r') as file:
            try:
                return yaml.safe_load(file)
            except yaml.YAMLError:
                return self._get_default_config()

    def _get_default_config(self):
        """Return default configuration"""
        return {
            "cors": {
                "origins": ["*"],
                "allow_credentials": True,
                "allow_methods": ["*"],
                "allow_headers": ["*"]
            },
            "servers": {
                "model_server": {
                    "url": "http://model-server:8000",
                    "timeout": 30
                },
                "inference_server": {
                    "url": "http://localhost:8000",
                    "timeout": 30
                },
                "inference_server_2": {
                    "url": "http://172.18.116.37:8000",
                    "timeout": 30
                }
            }
        }

    def get_config(self):
        """Return the loaded configuration"""
        return self.config